from FCN_LSTM_moudle import FCN_LSTM
from Get_DataSet import GetDataset
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import torch
import numpy as np

# 查看当前是否存在可用gpu
if torch.cuda.is_available():
    device = torch.device("cuda")  # 使用GPU
else:
    device = torch.device("cpu")  # 使用CPU

# 定义训练过程函数，batch_size：每次取得样本个数，train_times：训练次数，data_set_path：通过GetDataset函数得到的数据集，loss_list：损失函数输出值的存放列表
def Train_model(batch_size, model:FCN_LSTM, train_times, data_set_path, loss_list):
    
    # 损失函数为交叉熵
    criterion = nn.CrossEntropyLoss(weight = class_weights)
    criterion = criterion.to(device)
    # 优化器选择Adam
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    # 输出训练次数
    print("train times:", 0, end='', flush=True)  # 第一次输出，不换行
    best_loss = 1e8
    loss_nochanged = 0
    
    # 训练epoch次
    for epoch in range(train_times):   
        # 加载数据集
        train_loader = DataLoader(dataset=data_set_path, batch_size=batch_size, shuffle=True)
        
        # 分批次训练
        for train_data in train_loader:
            
            x_train = train_data[0]
            y_train = train_data[1]-1 #真实标签值的取值必须为[0,max_classes)
            
            optimizer.zero_grad() 
            # 将数据放到device上
            x_train = x_train.to(device)
            y_train = y_train.to(device)
            # 前向传播
            y_out = model.forward(x_train)
            # 交叉熵函数要求真实标签为long整形的一维张量，而模型输出结果为float形，因此需要修改y_train的数据类型
            y_train = y_train.squeeze()           
            y_train = y_train.long()
            
            # print(y_train.dtype)
            # print(y_out.dtype)
            # 损失函数计算
            loss = criterion(y_out,y_train)
            # 后向传播
            loss.backward()
            # 优化器优化
            optimizer.step()
            
            if best_loss > loss.item():
                best_loss = loss.item()
           
        if len(loss_list) > train_times//100:
            if loss.item() < loss_list[-1-train_times//100 +1]: #若1oss相比于之前1%次训练时下降了则保存模型和数据
                torch.save(model, 'LstmFcn_model.pth')
                torch.save(optimizer, 'LstmFcn_optimizer.pth') 
            
        if loss.item() <= best_loss * 1.1:
            loss_nochanged += 1 
        else:
            loss_nochanged=0
        
        if loss_nochanged >= (train_times*5)// 100: #连续5%的训练次数中loss无明显下降，则增大学习率
            for param_group in optimizer.param_groups:
                param_group['lr'] *= 10
            
        # 将loss加入loss列表
        loss_list.append(loss.item())  
        print('\r', end='')  # 使用\r实现回车，回到行的开头
        print("train times:", epoch, end='', flush=True)  # 仅更新数值部分

def Evaluate_model(batch_size,model:FCN_LSTM, data_set_path):
    predictions = []
    targets = []
    test_loader = DataLoader(dataset=data_set_path, batch_size=batch_size, shuffle=True)
    with torch.no_grad():
        for test_data in test_loader:
            x_test = test_data[0]
            y_test = test_data[1]-1 #真实标签值的取值必须为[0,max_classes)
            
            outputs = model(x_test)
            print('model outputs is', outputs.shape)
            print('pridect is', y_test.shape)
            _, predicted = torch.max(outputs.data, 1)  # 根据模型输出选择预测类别
            predictions.extend(predicted.cpu().numpy())
            targets.extend(y_test.cpu().numpy())

    return predictions, targets
    return
    
    

def draw_Loss(loss_list):
    plt.plot(loss_list)
    plt.xlabel('train times')
    plt.ylabel('loss')
    plt.title('loss wave')
    plt.show()
    
# 设置损失函数列表
loss_values = []
# 获得数据集
data_set = GetDataset('./data/ECG5000_TRAIN')

# 获得最大序列长度，用于构造模型
exp_data = DataLoader(dataset=data_set, batch_size=10, shuffle=True)
sequence_length = list(exp_data)[0][0].size(1)

#########################以下是获取每个类别标签的权重##############################################
labels_tensor = data_set.Get_labels().squeeze().long()-1
class_counts = torch.bincount(labels_tensor)
class_weights = (class_counts.float()) /  len(labels_tensor) 
# print(class_weights)
############################################################################################
# 构造模型并放入device
Model = FCN_LSTM(max_sequence_length = sequence_length, nb_class = data_set.Get_Classes(), num_cells = 8)
Model = Model.to(device)

# 训练模型
Model.train()#设置为训练模式
Train_model(batch_size=50, model=Model, train_times = 1000, data_set_path = data_set, loss_list = loss_values)

# 绘制loss随训练次数变化曲线
draw_Loss(loss_list = loss_values)


Test_data_set = GetDataset('./data/ECG5000_TRAIN')
Model.eval()#设置为评估模式
preds, targets = Evaluate_model(batch_size=10,model=Model,data_set_path=Test_data_set)
accuracy = np.mean(np.array(preds) == np.array(targets))
print(f"准确率: {accuracy}")

